package info.wangyuan.agent.service.rag;


import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.convert.Convert;
import cn.hutool.core.util.StrUtil;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.filter.comparison.IsIn;
import info.wangyuan.agent.common.Constant.EmbeddingConstant;
import info.wangyuan.agent.config.ModelFactory;
import info.wangyuan.agent.entity.dto.RagDocSyncStorageMsgDTO;
import info.wangyuan.agent.entity.po.DocumentUnit;
import info.wangyuan.agent.entity.po.FileInfo;
import info.wangyuan.agent.mapper.DocumentUnitMapper;
import info.wangyuan.agent.mapper.FileInfoMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.*;
import java.util.stream.Collectors;

import static dev.langchain4j.store.embedding.filter.MetadataFilterBuilder.metadataKey;

/**
 * 向量化存储
 *
 * @author Albert
 * @since 2025-08-31 23:28:31
 */
@Slf4j
@Component
public class RagEmbeddingDomainService {

    private final ModelFactory modelFactory;
    private final FileInfoMapper fileInfoMapper;
    private final RerankDomainService rerankService;
    private final DocumentUnitMapper documentUnitMapper;
    private final EmbeddingStore<TextSegment> embeddingStore;


    public RagEmbeddingDomainService(ModelFactory modelFactory, FileInfoMapper fileInfoMapper, RerankDomainService rerankService, DocumentUnitMapper documentUnitMapper, EmbeddingStore<TextSegment> embeddingStore) {
        this.modelFactory = modelFactory;
        this.fileInfoMapper = fileInfoMapper;
        this.rerankService = rerankService;
        this.documentUnitMapper = documentUnitMapper;
        this.embeddingStore = embeddingStore;
    }

    /**
     * RAG文档检索（支持高级参数和缓存优化）
     *
     * @param dataSetIds           知识库ID集合
     * @param question             查询问题
     * @param maxResults           最大返回结果数量
     * @param minScore             最小相似度阈值
     * @param enableRerank         是否启用重排序
     * @param candidateMultiplier  候选结果倍数
     * @param enableQueryExpansion 是否启用查询扩展（相邻片段）
     * @return 相关文档列表
     */
    public List<DocumentUnit> ragDoc(List<Integer> dataSetIds,
                                     String question,
                                     Integer maxResults,
                                     Double minScore,
                                     Boolean enableRerank,
                                     Integer candidateMultiplier,
                                     Boolean enableQueryExpansion) {

        // 参数验证
        if (CollUtil.isEmpty(dataSetIds)) {
            log.warn("知识库ID列表为空");
            return Collections.emptyList();
        }
        if (StrUtil.isBlank(question)) {
            log.warn("用户问题为空");
            return Collections.emptyList();
        }

        // 设置默认值和合理上限
        int finalMaxResults = Optional.ofNullable(maxResults).map(v -> Math.min(v, 100)).orElse(15);
        double finalMinScore = Optional.ofNullable(minScore).map(v -> Math.min(Math.max(v, 0.0), 1.0)).orElse(0.7);
        boolean finalEnableRerank = Optional.ofNullable(enableRerank).orElse(true);
        int finalCandidateMultiplier = Optional.ofNullable(candidateMultiplier).map(v -> Math.max(1, Math.min(v, 5))).orElse(2);

        // 记录搜索开始时间
        long startTime = System.currentTimeMillis();
        try {
            // 创建嵌入模型实例
            EmbeddingModel embeddingModel = modelFactory.getCustomEmbeddingModel();

            // 向量搜索 - 根据是否启用重排序决定搜索数量
            int searchLimit = finalEnableRerank
                    ? Math.max(finalMaxResults * finalCandidateMultiplier, 30)
                    : finalMaxResults;

            log.debug("开始 RAG 搜索: dataSetIds={}, question='{}', maxResults={}, minScore={}, enableRerank={}, searchLimit={}",
                    dataSetIds, question, finalMaxResults, finalMinScore, finalEnableRerank, searchLimit);

            // ===== 1. 向量检索 =====
            EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(EmbeddingSearchRequest.builder()
                    .filter(new IsIn(EmbeddingConstant.DATASET_ID, dataSetIds))
                    .maxResults(searchLimit)
                    .minScore(finalMinScore)
                    .queryEmbedding(Embedding.from(embeddingModel.embed(question).content().vector()))
                    .build());
            List<EmbeddingMatch<TextSegment>> embeddingMatches = searchResult.matches();

            // ===== 2. 重排序（可选） =====
            if (finalEnableRerank && CollUtil.isNotEmpty(embeddingMatches)) {
                long rerankStart = System.currentTimeMillis();
                embeddingMatches = rerankService.rerankDocument(searchResult, question);
                log.debug("应用重排序: {} 个匹配项, 耗时 {} ms", embeddingMatches.size(), System.currentTimeMillis() - rerankStart);
            } else {
                log.debug("跳过重排序, 使用 {} 个匹配结果", embeddingMatches.size());
            }

            // ===== 3. 回退搜索（如果没有找到相关文档，尝试降低相似度阈值再次搜索） =====
            if (embeddingMatches.isEmpty() && finalMinScore > 0.3) {
                log.info("未找到符合最低分数 {} 的结果, 降低相似度阈值至 0.3 重试", finalMinScore);
                EmbeddingSearchResult<TextSegment> fallbackResult = embeddingStore.search(EmbeddingSearchRequest.builder()
                        .filter(new IsIn(EmbeddingConstant.DATASET_ID, dataSetIds))
                        .maxResults(searchLimit)
                        .minScore(0.3)
                        .queryEmbedding(Embedding.from(embeddingModel.embed(question).content().vector()))
                        .build());
                embeddingMatches = fallbackResult.matches();
                log.debug("回退搜索找到 {} 个匹配项", embeddingMatches.size());
            }

            // ===== 4.提取文档ID并创建ID到分数的映射 =====
            Map<Integer, Double> documentScores = new HashMap<>();
            List<Integer> documentIds = embeddingMatches.stream()
                    .limit(finalMaxResults)
                    .map(match -> {
                        Integer docId = match.embedded().metadata().getInteger(EmbeddingConstant.DOCUMENT_ID);
                        if (docId != null) {
                            documentScores.put(docId, match.score());
                            log.debug("找到文档：{}，得分：{}", docId, String.format("%.4f", match.score()));
                        }
                        return docId;
                    })
                    .filter(Objects::nonNull)
                    .toList();

            if (CollUtil.isEmpty(documentIds)) {
                log.info("查询 '{}' 未找到相关文档，最低分数={}", question, finalMinScore);
                return Collections.emptyList();
            }

            // ===== 5. 查询扩展：如果启用了查询扩展，添加相邻片段 =====
            List<Integer> finalDocIds = new ArrayList<>(documentIds);
            if (Boolean.TRUE.equals(enableQueryExpansion)) {
                List<DocumentUnit> initialDocs = documentUnitMapper.selectList(
                        Wrappers.lambdaQuery(DocumentUnit.class)
                                .in(DocumentUnit::getId, documentIds)
                );

                // 收集所有需要的片段ID（使用LinkedHashSet保持顺序并去重）
                Set<Integer> expandedIds = new LinkedHashSet<>(documentIds);
                for (DocumentUnit doc : initialDocs) {
                    // 查询相邻页面片段（前一页、当前页、后一页）
                    List<DocumentUnit> adjacent = documentUnitMapper.selectList(
                            Wrappers.lambdaQuery(DocumentUnit.class)
                                    .eq(DocumentUnit::getFileId, doc.getFileId())
                                    .between(DocumentUnit::getPageIndex, Math.max(1, doc.getPageIndex() - 1), doc.getPageIndex() + 1)
                                    .eq(DocumentUnit::getIsVectored, true)
                    );
                    adjacent.forEach(chunk -> expandedIds.add(chunk.getId()));
                }
                finalDocIds = new ArrayList<>(expandedIds);
                log.info("查询扩展启用：原始 {} 个 chunk 扩展为 {} 个 chunk", documentIds.size(), finalDocIds.size());
            }

            // ===== 6. 查询数据库并组装结果 =====
            List<DocumentUnit> allDocs = documentUnitMapper.selectList(
                    Wrappers.lambdaQuery(DocumentUnit.class)
                            .in(DocumentUnit::getId, finalDocIds)
            );

            // 转为 Map，加速后续匹配
            Map<Integer, DocumentUnit> allDocsMap = allDocs.stream()
                    .collect(Collectors.toMap(DocumentUnit::getId, d -> d));

            // 按照检索相关性顺序重新排列结果，并设置相似度分数
            List<DocumentUnit> sortedResults = finalDocIds.stream()
                    .map(id -> {
                        DocumentUnit doc = allDocsMap.get(id);
                        if (doc != null) {
                            // 设置相似度分数：原始匹配使用向量搜索分数，扩展片段使用默认分数
                            double score = documentScores.getOrDefault(id, finalMinScore * 0.8);
                            doc.setSimilarityScore(score);
                        }
                        return doc;
                    })
                    .filter(Objects::nonNull)
                    .toList();

            // ===== 7. 记录搜索性能统计 =====
            long totalTime = System.currentTimeMillis() - startTime;
            double avgScore = embeddingMatches.stream().mapToDouble(EmbeddingMatch::score).average().orElse(0.0);

            log.info("RAG 搜索完成: query='{}', 返回 {} 个文档, 平均得分={}, 耗时={} ms",
                    question, sortedResults.size(), String.format("%.4f", avgScore), totalTime);

            return sortedResults;

        } catch (Exception e) {
            log.error("RAG 搜索异常: query='{}', 耗时={} ms", question, System.currentTimeMillis() - startTime, e);
            return Collections.emptyList();
        }
    }

    /**
     * 批量向量化入库
     */
    public void syncEmbeddingStorageByBatch(List<RagDocSyncStorageMsgDTO> ragDocSyncStorageMsgDTOList) {
        ragDocSyncStorageMsgDTOList.forEach(this::syncEmbeddingStorage);
    }

    /**
     * 文本向量化入库
     */
    public void syncEmbeddingStorage(RagDocSyncStorageMsgDTO ragDocSyncStorageMsgDTO) {
        final Integer docId = ragDocSyncStorageMsgDTO.getId();

        final DocumentUnit documentUnit = documentUnitMapper.selectById(docId);
        if (documentUnit == null) return;

        final String content = documentUnit.getContent();

        final Metadata documentMetadata = buildMetadata(ragDocSyncStorageMsgDTO);

        final TextSegment textSegment = new TextSegment(content, documentMetadata);

        // 使用消息中配置的嵌入模型生成向量
        var embeddingModel = ragDocSyncStorageMsgDTO.getEmbeddingModel();
        Embedding embeddings = embeddingModel.embed(textSegment).content();

        embeddingStore.add(embeddings, textSegment);

        documentUnitMapper.update(
                Wrappers.lambdaUpdate(DocumentUnit.class)
                        .eq(DocumentUnit::getId, docId)
                        .set(DocumentUnit::getIsVectored, true)
        );

        // 修改文件状态
        final FileInfo fileInfo = fileInfoMapper.selectById(ragDocSyncStorageMsgDTO.getFileId());

        final Integer pageSize = fileInfo.getPageSize();

        final Long isVector = documentUnitMapper.selectCount(
                Wrappers.lambdaQuery(DocumentUnit.class)
                        .eq(DocumentUnit::getFileId, documentUnit.getFileId())
                        .eq(DocumentUnit::getIsVectored, true)
        );

        final Integer anInt = Convert.toInt(isVector);

        if (anInt >= pageSize) {
            // 文件向量化完成
            log.info("文件 {} 第 {} 页向量化完成", documentUnit.getFileId(), documentUnit.getPageIndex());
        }
    }

    private Metadata buildMetadata(RagDocSyncStorageMsgDTO ragDocSyncStorageMsgDTO) {
        final Metadata metadata = new Metadata();
        metadata.put(EmbeddingConstant.FILE_ID, ragDocSyncStorageMsgDTO.getFileId());
        metadata.put(EmbeddingConstant.FILE_NAME, ragDocSyncStorageMsgDTO.getFileName());
        metadata.put(EmbeddingConstant.DOCUMENT_ID, ragDocSyncStorageMsgDTO.getId());
        metadata.put(EmbeddingConstant.DATASET_ID, ragDocSyncStorageMsgDTO.getDatasetId());
        return metadata;
    }

    /**
     * 批量删除文件向量
     *
     * @param fileIds 文件id集合
     */
    public void removeEmbeddingByFileIds(List<Integer> fileIds) {
        if (CollUtil.isEmpty(fileIds)) return;
        embeddingStore.removeAll(metadataKey(EmbeddingConstant.FILE_ID).isIn(fileIds));
    }

    /**
     * 删除单个文件向量
     *
     * @param fileId 文件ID
     */
    public void removeEmbeddingByFileId(Integer fileId) {
        if (fileId == null) return;
        removeEmbeddingByFileIds(Collections.singletonList(fileId));
    }

}
